#!/usr/bin/env python3

# Copyright 2023 Kry10 Limited
# SPDX-License-Identifier: BSD-2-Clause

# Export kernel build artifacts for binary verification. This includes
# kernel binaries, disassembly, and if present, the C graph-lang model
# extracted from the Isabelle C spec.

# GitHub l4v proof workflows use this to save kernel build outputs as
# workflow artifacts. The artifacts are then picked up by binary
# verification workflows. This ensures that the kernel builds used in
# binary verification are generated by the same toolchains used in the
# Isabelle proofs.

# For binary verification, we perform builds at two optimisation levels.
# We also save a copy of the preprocessed source used in the Isabelle
# proofs, so we can later compare it to the preprocessed source used in
# binary verification.

# We expect the L4V_ARCH environment variable to be set as usual. We use
# this to locate kernel sources built by the l4v proofs, and the C
# graph-lang file. Any other environment variables required by the l4v
# kernel build system are assumed to be already set.

# We require a single command-line argument, which is the directory into
# which builds will be exported. Under this directory, builds will be
# placed into subdirectories according to their L4V_ARCH, L4V_FEATURES,
# and optimisation level. With L4V_ARCH=RISCV64 and L4V_FEATURES=MCS, for
# example:

# RISCV64-MCS-O1/
# RISCV64-MCS-O1/kernel_all.cpp
# RISCV64-MCS-O1/kernel.elf
# RISCV64-MCS-O1/...
# RISCV64-MCS-O2/
# RISCV64-MCS-O2/kernel_all.cpp
# RISCV64-MCS-O2/kernel.elf
# RISCV64-MCS-O1/...

import argparse
import os
import shutil
import subprocess
import sys

from pathlib import Path
from tempfile import TemporaryDirectory
from typing import List, NamedTuple, Optional


class L4vPaths(NamedTuple):
    kernel_mk: Path
    c_pp: Path
    c_functions: Path


def get_l4v_paths(l4v_arch: str) -> L4vPaths:
    l4v_c = Path(__file__).resolve().parent
    l4v = l4v_c.parent.parent.parent
    kernel_mk = l4v_c / 'kernel.mk'
    c_pp = l4v_c / 'build' / l4v_arch / 'kernel_all.c_pp'
    c_functions = l4v / 'proof' / 'asmrefine' / 'export' / l4v_arch / 'CFunDump.txt'
    return L4vPaths(kernel_mk=kernel_mk, c_pp=c_pp, c_functions=c_functions)


def path_suffix(opt_suffix: Optional[str]) -> str:
    return f'-{opt_suffix}' if opt_suffix else ''


class ExportConfig(NamedTuple):
    export_root: Path
    l4v_arch: str
    l4v_features: Optional[str]
    l4v_plat: Optional[str]
    l4v_paths: L4vPaths
    manifest: Optional[Path]

    def config_name(self, optimisation: str) -> str:
        features = path_suffix(self.l4v_features)
        plat = path_suffix(self.l4v_plat)
        return f'{self.l4v_arch}{features}{plat}{optimisation}'

    def export_path(self, optimisation: str) -> Path:
        return self.export_root / self.config_name(optimisation)

    def do_export(self, optimisation: str) -> None:
        config_name = self.config_name(optimisation)
        export_dir = self.export_path(optimisation)

        print(f'Exporting kernel build for {config_name} to {export_dir}...')

        if not self.l4v_paths.c_pp.is_file():
            print('  Note: No l4v kernel build found.')
        if not self.l4v_paths.c_functions.is_file():
            print('  Note: No C graph-lang found.')

        with TemporaryDirectory() as build_tmp:
            env = {
                **os.environ,
                'KERNEL_BUILD_ROOT': Path(build_tmp) / 'build',
                'KERNEL_EXPORT_DIR': export_dir,
                'CONFIG_OPTIMISATION': optimisation
            }
            p = subprocess.run(['make', '-f', self.l4v_paths.kernel_mk, 'kernel_export'],
                               env=env, stdin=subprocess.DEVNULL)
        if p.returncode != 0:
            print(f'Error: Kernel build for {config_name} failed', file=sys.stderr)
            sys.exit(1)

        # Copy the preprocessed source from the l4v build, so we can later compare
        # it to the exported builds if binary verification fails. There's no point
        # comparing here, because there will always be differences due to lines
        # inserted by the preprocessor and config system.
        if self.l4v_paths.c_pp.is_file():
            shutil.copyfile(self.l4v_paths.c_pp, export_dir / 'kernel_all.c_pp.l4v')

        # Copy the C graph-lang, if it exists. Note that C graph-lang might not
        # have been produced, if SimplExportAndRefine is not enabled for this
        # configuration, or if proofs were cached. It's therefore not an error for
        # it not to be present.
        if self.l4v_paths.c_functions.is_file():
            shutil.copyfile(self.l4v_paths.c_functions, export_dir / 'CFunctions.txt')

        # If a manifest was given on the command line, copy that to the output.
        if self.manifest is not None:
            shutil.copyfile(self.manifest, export_dir / 'manifest.xml')

        # Write a file describing the configuration.
        with open(export_dir / 'config.env', 'w') as config_env:
            config_env.write(f'L4V_ARCH={self.l4v_arch}\n')
            config_env.write(f'L4V_FEATURES={self.l4v_features or ""}\n')
            config_env.write(f'L4V_PLAT={self.l4v_plat or ""}\n')
            config_env.write(f'CONFIG_OPTIMISATION={optimisation}\n')


class ExportCommand(NamedTuple):
    config: ExportConfig
    optimisations: List[str]
    force: bool

    def run(self) -> None:
        if not self.config.l4v_paths.c_functions.is_file() and not self.force:
            print('Will not export kernel builds, because no C graph-lang was found.')
            return
        for optimisation in self.optimisations:
            self.config.do_export(optimisation)


def parse_args() -> ExportCommand:
    parser = argparse.ArgumentParser(
        description='Export kernel build artifacts.')
    parser.add_argument('--export-root', metavar='DIRECTORY', type=Path, required=True,
                        help='Export directory')
    parser.add_argument('--manifest-xml', metavar='FILENAME', type=Path,
                        help='Repo manifest xml file')
    parser.add_argument('-O', metavar='OPTIMISATION_LEVEL', type=str, nargs=1,
                        choices=['1', '2'], dest='optimisations',
                        help='Optimisation level')
    parser.add_argument('--force', dest='force', action='store_true',
                        help='Build even when no C graph-lang is found')
    parser.set_defaults(optimisations=['1', '2'], force=False)
    args = parser.parse_args()

    l4v_arch = os.environ.get('L4V_ARCH')
    if not l4v_arch:
        print('error: L4V_ARCH not set', file=sys.stderr)
        sys.exit(1)

    if args.manifest_xml and not args.manifest_xml.is_file():
        print(f'error: bad manifest: {self.manifest_xml}', file=sys.stderr)
        sys.exit(1)

    config = ExportConfig(export_root=args.export_root,
                          l4v_arch=l4v_arch,
                          l4v_features=os.environ.get('L4V_FEATURES'),
                          l4v_plat=os.environ.get('L4V_PLAT'),
                          l4v_paths=get_l4v_paths(l4v_arch),
                          manifest=args.manifest_xml)

    optimisations = (f'-O{o}' for o in args.optimisations)
    return ExportCommand(config=config,
                         optimisations=optimisations,
                         force=args.force)


def main(cmd: ExportCommand) -> int:
    cmd.run()
    return 0


if __name__ == '__main__':
    exit(main(parse_args()))
